import sqlite3  # 新增：导入 SQLite 数据库模块

from PyQt5 import QtGui  # 修改：正确导入 QtGui 模块
from PyQt5.QtCore import QThread, pyqtSignal
from PyQt5.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QWidget, QTextEdit, QLineEdit, QPushButton, \
    QHBoxLayout, QDialog, QMessageBox, QLabel
from openai import OpenAI


def get_api_key():
    """从数据库中获取 api_key"""
    conn = sqlite3.connect('data/questions.db')
    cursor = conn.cursor()
    cursor.execute('SELECT api_key FROM settings LIMIT 1')
    result = cursor.fetchone()
    conn.close()
    return result[0] if result else None

def set_api_key(api_key):
    """将 api_key 存储到数据库"""
    conn = sqlite3.connect('data/questions.db')
    cursor = conn.cursor()
    cursor.execute('INSERT OR REPLACE INTO settings (id, api_key) VALUES (1, ?)', (api_key,))
    conn.commit()
    conn.close()



# 修改：从数据库中读取 api_key
api_key =""
if api_key:
    client = OpenAI(api_key=api_key, base_url="https://api.siliconflow.cn/v1")
else:
    client = None  # 如果没有 api_key，则不初始化客户端

# 配置参数
max_messages = 10  # messages 的最大容量

# 存储聊天记录
messages = []

class WorkerThread(QThread):
    """新增：使用 QThread 处理 OpenAI 请求"""
    update_chat_signal = pyqtSignal(str)  # 用于更新聊天记录的信号
    error_signal = pyqtSignal(str)  # 用于传递错误信息的信号

    def __init__(self, user_message):
        super().__init__()
        self.user_message = user_message

    def run(self):
        try:
            response = client.chat.completions.create(
                model="Qwen/Qwen2.5-72B-Instruct",
                messages=messages,
                stream=True
            )

            assistant_response = ""
            current_line = ""  # 新增：用于累积当前行的内容

            for chunk in response:
                if not chunk.choices:
                    continue

                # 处理 content 字段
                if chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    lines = content.split("\n")  # 按换行符分割内容
                    for i, line in enumerate(lines):
                        if i == 0 and current_line:  # 如果是第一行且有累积内容
                            current_line += line
                        else:
                            current_line = line  # 否则重置当前行
                        if i < len(lines) - 1 or "\n" in content:  # 如果是中间行或末尾有换行符
                            self.update_chat_signal.emit(f"<p>{current_line}</p>")  # 发送信号更新界面
                            current_line = ""  # 清空当前行
                    assistant_response += content

                # 处理 reasoning_content 字段
                if chunk.choices[0].delta.reasoning_content:
                    reasoning_content = chunk.choices[0].delta.reasoning_content
                    lines = reasoning_content.split("\n")  # 按换行符分割内容
                    for i, line in enumerate(lines):
                        if i == 0 and current_line:  # 如果是第一行且有累积内容
                            current_line += line
                        else:
                            current_line = line  # 否则重置当前行
                        if i < len(lines) - 1 or "\n" in reasoning_content:  # 如果是中间行或末尾有换行符
                            self.update_chat_signal.emit(f"<p>{current_line}</p>")  # 发送信号更新界面
                            current_line = ""  # 清空当前行
                    assistant_response += reasoning_content

            # 如果最后还有未处理的当前行内容
            if current_line:
                self.update_chat_signal.emit(f"<p>{current_line}</p>")

            # 添加助手消息到 messages
            messages.append({'role': 'assistant', 'content': assistant_response})
            if len(messages) > max_messages:
                messages.pop(0)  # 移除最早的记录

        except Exception as e:
            # 捕获异常并通过信号传递错误信息
            self.error_signal.emit(f"<p><strong>错误:</strong> {str(e)}</p> <strong style='color:red;'>(请确认是否正确设置 API KEY)</strong>")

class ChatBotWindow(QMainWindow):
    def __init__(self,main_window):
        super().__init__()
        api_key = get_api_key()
        self.main_window = main_window  # 保存main_window引用
        self.init_ui()
        self.is_requesting = False

    def init_ui(self):
        self.setWindowTitle("聊天机器人")
        self.setGeometry(100, 100, 600, 400)

        # 创建主布局
        main_layout = QVBoxLayout()

        # 聊天记录显示区域
        self.chat_display = QTextEdit(self)
        self.chat_display.setReadOnly(True)
        self.chat_display.setHtml("<html><body></body></html>")  # 修改：初始化为HTML格式
        main_layout.addWidget(self.chat_display)

        # 输入框和按钮布局
        input_layout = QHBoxLayout()

        # 输入框
        self.input_box = QLineEdit(self)
        self.input_box.returnPressed.connect(self.send_message)
        input_layout.addWidget(self.input_box)

        # 发送按钮
        self.send_button = QPushButton("发送", self)
        self.send_button.clicked.connect(self.send_message)
        input_layout.addWidget(self.send_button)

        # 断开链接按钮
        self.disconnect_button = QPushButton("断开链接", self)
        self.disconnect_button.clicked.connect(self.disconnect)
        input_layout.addWidget(self.disconnect_button)

        # 设置按钮
        self.settings_button = QPushButton("设置 ⚙️", self)
        self.settings_button.clicked.connect(self.open_settings)
        input_layout.addWidget(self.settings_button)

        main_layout.addLayout(input_layout)

        # 设置中央部件
        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)

    def send_message(self):
        if self.is_requesting:  # 如果正在请求中，禁止发送
            return

        user_message = self.input_box.text().strip()
        if not user_message:
            return

        self.is_requesting = True  # 标记为正在请求中
        self.send_button.setText("思考中...🤔")  # 修改：显示“思考中...”
        self.send_button.setEnabled(False)  # 禁用发送按钮
        self.input_box.setEnabled(False)  # 禁用输入框

        # 添加用户消息到 messages
        messages.append({'role': 'user', 'content': user_message})
        if len(messages) > max_messages:
            messages.pop(0)  # 移除最早的记录

        # 显示用户消息
        self.append_to_chat(f"<p><strong>小主:</strong> {user_message}</p>")

        # 使用 QThread 处理 OpenAI API 请求
        self.worker_thread = WorkerThread(user_message)
        self.worker_thread.update_chat_signal.connect(self.append_to_chat)
        self.worker_thread.error_signal.connect(self.handle_error)
        self.worker_thread.finished.connect(self.on_request_finished)
        self.worker_thread.start()

    def append_to_chat(self, markdown):
        """新增：将Markdown内容追加到聊天记录中"""
        current_html = self.chat_display.toHtml()
        # 检查是否为用户消息
        if "<strong>小主:</strong>" in markdown:
            # 将用户消息设置为绿色加粗
            html_content = markdown.replace("<strong>小主:</strong>", "<strong style='color:green;'>小主:</strong>")
        else:
            # 将Markdown内容转换为HTML并追加到现有内容中
            html_content = markdown.replace("**", "<strong>").replace("**", "</strong>")
            html_content = html_content.replace("##### ", "<h5>").replace("#####", "</h5>")
            html_content = html_content.replace("#### ", "<h4>").replace("####", "</h4>")
            html_content = html_content.replace("### ", "<h3>").replace("###", "</h3>")
            html_content = html_content.replace("## ", "<h2>").replace("##", "</h2>")
            html_content = html_content.replace("# ", "<h1>").replace("#", "</h1>")
        new_html = current_html.replace("</body>", f"<p>{html_content}</p></body>")
        self.chat_display.setHtml(new_html)
        # 新增：将滚动条移动到最底部
        self.chat_display.moveCursor(QtGui.QTextCursor.End)  # 修改：使用正确的 QtGui.QTextCursor

    def handle_error(self, error_message):
        """新增：处理错误信息"""
        self.append_to_chat(error_message)

    def on_request_finished(self):
        """新增：请求完成后恢复 UI 状态"""
        self.is_requesting = False  # 标记为请求结束
        self.send_button.setText("发送")  # 修改：恢复发送按钮文本
        self.send_button.setEnabled(True)  # 重新启用发送按钮
        self.input_box.setEnabled(True)  # 重新启用输入框
        self.input_box.clear()  # 清空输入框内容
        self.input_box.setFocus()  # 新增：将焦点设置到输入框

    def open_settings(self):
        """打开设置界面"""
        settings_window = SettingsWindow(self)
        settings_window.exec_()

    def update_api_key(self, api_key):
        """更新 api_key 并重新初始化客户端"""
        global client
        set_api_key(api_key)
        client = OpenAI(api_key=api_key, base_url="https://api.siliconflow.cn/v1")

    def disconnect(self):
        """新增：断开链接功能"""
        # 这里可以添加断开链接的具体逻辑
        self.chat_display.append("<p><strong>链接已断开</strong></p>")

# 新增：设置界面类
class SettingsWindow(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("设置")
        self.setGeometry(200, 200, 300, 200)  # 增加窗口高度以容纳更多内容

        layout = QVBoxLayout()

        # 提示信息
        self.info_label = QLabel("请到以下地址获取 API Key：", self)
        layout.addWidget(self.info_label)

        # 跳转链接
        self.link_label = QLabel('<a href="https://cloud.siliconflow.cn/account/ak">https://cloud.siliconflow.cn/account/ak</a>', self)
        self.link_label.setOpenExternalLinks(True)  # 允许点击跳转
        layout.addWidget(self.link_label)

        # API Key 输入框
        self.api_key_input = QLineEdit(self)
        self.api_key_input.setPlaceholderText("请输入 API Key")
        layout.addWidget(self.api_key_input)

        # 保存按钮
        save_button = QPushButton("保存", self)
        save_button.clicked.connect(self.save_api_key)
        layout.addWidget(save_button)

        self.setLayout(layout)

        # 居中显示窗口
        self.center_on_screen()

    def center_on_screen(self):
        """将窗口居中显示"""
        screen_geometry = QApplication.desktop().screenGeometry()
        x = (screen_geometry.width() - self.width()) // 2
        y = (screen_geometry.height() - self.height()) // 2
        self.move(x, y)

    def save_api_key(self):
        """保存 API Key 并关闭窗口"""
        api_key = self.api_key_input.text().strip()
        if api_key:
            self.parent().update_api_key(api_key)
            self.accept()
        else:
            QMessageBox.warning(self, "警告", "API Key 不能为空！")

# if __name__ == "__main__":
#     app = QApplication(sys.argv)
#     window = ChatBotWindow()
#     window.show()
#     sys.exit(app.exec())